# 安装必要库 (确保您已安装 modelscope kokoro pyaudio numpy librosa)
# pip install modelscope kokoro>=0.9.4 pyaudio numpy librosa

import torch
import numpy as np
import pyaudio
import time
import librosa
from kokoro import KPipeline

class KokoroNoiseReducedStreamer:
    """
    灵活的 Kokoro 流式合成器。
    支持可选的基于频谱的降噪和音量归一化，以提高播放质量。
    """
    def __init__(self, lang_code='z', voice='af_heart', nr_strength=0.0, sample_rate=24000):
        """
        初始化流式合成器和 PyAudio。
        :param lang_code: 语言代码，'z' 表示中文。
        :param voice: 声音名称，例如 'af_heart'。
        :param nr_strength: 降噪强度 (0.0 到 1.0)。设置为 0.0 则关闭降噪。
        """
        self.pipeline = KPipeline(lang_code=lang_code)
        self.voice = voice
        self.sample_rate = sample_rate
        self.nr_strength = nr_strength
        self.normalization_factor = 0.8  # 归一化目标最大值，避免削波
        
        # 初始化 PyAudio
        try:
            self.p = pyaudio.PyAudio()
        except Exception as e:
            print(f"错误: PyAudio 初始化失败: {e}")
            raise
        
    def _normalize_volume(self, audio_np):
        """将音频数据归一化到指定的最大值 (self.normalization_factor)"""
        max_val = np.max(np.abs(audio_np))
        if max_val > 0:
            audio_np = audio_np * (self.normalization_factor / max_val)
        return audio_np

    def _spectral_noise_reduction(self, audio_np):
        """基于频谱的降噪方法（谱减法）"""
        try:
            # 确保音频是浮点数格式
            audio_np = audio_np.astype(np.float32)

            # 使用 librosa 进行短时傅里叶变换 (STFT)
            n_fft, hop_length = 2048, 512
            stft = librosa.stft(audio_np, n_fft=n_fft, hop_length=hop_length)
            magnitude = np.abs(stft)
            phase = np.angle(stft)
            
            # 简单的噪声估计：取前几帧的平均幅度作为噪声基线
            num_noise_frames = min(magnitude.shape[1], 5)
            noise_estimate = np.mean(magnitude[:, :num_noise_frames], axis=1, keepdims=True)
            
            # 谱减法：从信号幅度中减去噪声估计，并应用地板值
            clean_magnitude = np.maximum(
                magnitude - self.nr_strength * noise_estimate, 
                0.001 * magnitude  # 设置一个较低的噪声地板，防止声音过度失真
            )
            
            # 重建信号并进行逆短时傅里叶变换 (ISTFT)
            clean_stft = clean_magnitude * np.exp(1j * phase)
            clean_audio = librosa.istft(clean_stft, hop_length=hop_length, length=len(audio_np))
            
            return clean_audio
            
        except Exception as e:
            print(f"警告: 频谱降噪处理失败: {e}")
            return audio_np

    def _preprocess_audio(self, audio_data):
        """音频预处理管道：降噪和音量归一化"""
        if torch.is_tensor(audio_data):
            audio_np = audio_data.numpy()
        else:
            audio_np = audio_data
        
        # 确保单声道且为 float32
        if len(audio_np.shape) > 1:
            audio_np = audio_np.flatten()
        audio_np = audio_np.astype(np.float32)

        # 1. 降噪
        if self.nr_strength > 0.0:
            audio_np = self._spectral_noise_reduction(audio_np)
        
        # 2. 音量归一化
        audio_np = self._normalize_volume(audio_np)
        
        return audio_np
    
    def _play_audio_chunk(self, audio_np):
        """使用 PyAudio 播放处理后的音频片段"""
        # 以浮点数格式打开音频流
        stream = self.p.open(
            format=pyaudio.paFloat32,
            channels=1,
            rate=self.sample_rate,
            output=True,
            frames_per_buffer=4096, # 适当增大缓冲区
        )
        
        try:
            # 直接写入整个音频数据。在默认的阻塞模式下，写入操作会等待直到数据全部被接受到音频缓冲区。
            # 移除了 while stream.is_active() 循环，避免不必要的阻塞和卡顿。
            stream.write(audio_np.tobytes())
            
        except Exception as e:
            print(f"错误: 音频播放失败: {e}")
        finally:
            # 停止流并关闭资源
            stream.stop_stream()
            stream.close()

    def stream_synthesis(self, text, chunk_delay=0):
        """流式合成文本并实时播放"""
        
        generator = self.pipeline(
            text, 
            voice=self.voice,
            speed=1.0,
            split_pattern=r'[。！？!?;\n]+' # 根据标点符号分割
        )
        
        print(f"开始流式合成 (降噪强度: {self.nr_strength})...")
        
        for i, (graphemes, _, audio) in enumerate(generator):
            print(f"\n🔊 播放片段 {i+1}: {graphemes}")
            
            # 音频预处理（降噪和归一化）
            processed_audio = self._preprocess_audio(audio)
            
            # 播放音频
            self._play_audio_chunk(processed_audio)
            
            # 延迟控制，让合成和播放更平滑
            # time.sleep(chunk_delay)
        
        print("\n✅ 流式合成完成！")
    
    def close(self):
        """释放 PyAudio 资源"""
        if hasattr(self, 'p'):
            self.p.terminate()

# 使用示例
if __name__ == "__main__":
    # 示例文本
    test_text = """您好，我已经收到了您的反馈，并理解了您遇到的两个问题：PyAudio 播放卡顿（KeyboardInterrupt）以及希望代码只保留降噪功能。

卡顿问题分析与修复：

您报告的卡顿现象 (while stream.is_active(): time.sleep(0.01)) 是 PyAudio 阻塞模式下常见的陷阱。该循环用于等待音频播放完成，但在某些系统或配置中，is_active() 的状态刷新可能不及时或不准确，导致程序长时间阻塞。

我已修改 _play_audio_chunk 函数，移除了这个不稳定的等待循环。现在，代码将依靠 PyAudio stream.write() 在阻塞模式下自身等待数据写入完成的机制来确保音频片段播放。

关于保留降噪功能：

KokoroNoiseReducedStreamer 类目前的设计已经满足了您的需求：

它是专门为降噪设计的流式合成器。

您可以通过设置 nr_strength=0.0 来关闭降噪（作为基准对比），或设置 0.1 到 1.0 的值来启用降噪。

我将更新 Canvas 中的代码，主要针对卡顿问题进行修复。

我已经修复了 Canvas 文件 kokoro_streamer_optimized.py 中的 PyAudio 播放阻塞问题。现在流式播放应该会更加流畅，不会再出现卡在播放循环中的情况。

您可以尝试运行修改后的代码，并观察降噪效果。如果想进一步调整降噪的强度，可以直接修改 nr_strength 的值。"""

    print("\n--- 演示 2: 开启降噪 (强度 0.4) ---")
    streamer_advanced = KokoroNoiseReducedStreamer(nr_strength=0.4)
    try:
        streamer_advanced.stream_synthesis(test_text)
    except Exception as e:
        print(f"程序运行错误: {e}")
    finally:
        streamer_advanced.close()
